/*
 * $Id$
 *
 *                This source code is part of
 *
 *                 G   R   O   M   A   C   S
 *
 *          GROningen MAchine for Chemical Simulations
 *
 *                        VERSION 3.0
 *
 * Copyright (c) 1991-2001
 * BIOSON Research Institute, Dept. of Biophysical Chemistry
 * University of Groningen, The Netherlands
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * If you want to redistribute modifications, please consider that
 * scientific software is very special. Version control is crucial -
 * bugs must be traceable. We will be happy to consider code for
 * inclusion in the official distribution, but derived work must not
 * be called official GROMACS. Details are found in the README & COPYING
 * files - if they are missing, get the official version at www.gromacs.org.
 *
 * To help us fund GROMACS development, we humbly ask that you cite
 * the papers on the package - you can find them in the top README file.
 *
 * Do check out http://www.gromacs.org , or mail us at gromacs@gromacs.org .
 *
 * And Hey:
 * Gyas ROwers Mature At Cryogenic Speed
 */


/* This line is only for CVS version info */


#include "statutil.h"
#include "typedefs.h"
#include "smalloc.h"
#include "vec.h"
#include "copyrite.h"
#include "statutil.h"
#include "tpxio.h"
#include "fftgrid.h"
#include "index.h"
#include "pbc.h"
#include "GetCharges.h"
#include "centerofmass.h"
#include "rdf_grid.h"
#include "get_potential.h"
#include "xvgr.h"
#include "get_xplor.h"
rvec ngrid={1,1,1};

real*** Rho_sol;
real*** Rho_slt;
real*** Rho_tot;
real*** Phi_sol;
real*** Phi_slt;
real*** Phi_tot;
real* Charges;
enum { euSel,euRect, euTric, euCompact, euNR};

static void calc_recipbox(matrix box,matrix recipbox)
{
  /* Save some time by assuming upper right part is zero */

  real tmp=1.0/(box[XX][XX]*box[YY][YY]*box[ZZ][ZZ]);

  recipbox[XX][XX]=box[YY][YY]*box[ZZ][ZZ]*tmp;
  recipbox[XX][YY]=0;
  recipbox[XX][ZZ]=0;
  recipbox[YY][XX]=-box[YY][XX]*box[ZZ][ZZ]*tmp;
  recipbox[YY][YY]=box[XX][XX]*box[ZZ][ZZ]*tmp;
  recipbox[YY][ZZ]=0;
  recipbox[ZZ][XX]=(box[YY][XX]*box[ZZ][YY]-box[YY][YY]*box[ZZ][XX])*tmp;
  recipbox[ZZ][YY]=-box[ZZ][YY]*box[XX][XX]*tmp;
  recipbox[ZZ][ZZ]=box[XX][XX]*box[YY][YY]*tmp;
}

static void put_molecule_com_in_box(int unitcell_enum,int ecenter,
                                    t_block *mols,
                                    int natoms,t_atom atom[],
                                    int ePBC,matrix box,rvec x[])
{
    atom_id i,j;
    int     d;
    rvec    com,new_com,shift,box_center;
    real    m;
    double  mtot;
    t_pbc   pbc;

    calc_box_center(ecenter,box,box_center);
    set_pbc(&pbc,ePBC,box);
    if (mols->nr <= 0)
        gmx_fatal(FARGS,"There are no molecule descriptions. I need a tpr file for this pbc option.");
    for(i=0; (i<mols->nr); i++) {
        /* calc COM */
        clear_rvec(com);
        mtot = 0;
        for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
            m = atom[j].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[j][d];
            mtot += m;
        }
        /* calculate final COM */
        svmul(1.0/mtot, com, com);

        /* check if COM is outside box */
        copy_rvec(com,new_com);
        switch (unitcell_enum) {
        case euRect:
            put_atoms_in_box(box,1,&new_com);
            break;
        case euTric:
            put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
            break;
        case euCompact:
            put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
            break;
        }
        rvec_sub(new_com,com,shift);
        if (norm2(shift) > 0) {
            if (debug)
                fprintf (debug,"\nShifting position of molecule %d "
                         "by %8.3f  %8.3f  %8.3f\n", i+1, PR_VEC(shift));
            for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
                rvec_inc(x[j],shift);
            }
        }
    }
}

static void put_residue_com_in_box(int unitcell_enum,int ecenter,
                                   int natoms,t_atom atom[],
                                   int ePBC,matrix box,rvec x[])
{
    atom_id i, j, res_start, res_end, res_nat;
    int     d, presnr;
    real    m;
    double  mtot;
    rvec    box_center,com,new_com,shift;

    calc_box_center(ecenter,box,box_center);

    presnr = NOTSET;
    res_start = 0;
    clear_rvec(com);
    mtot = 0;
    for(i=0; i<natoms+1; i++) {
        if (i == natoms || (presnr != atom[i].resind && presnr != NOTSET)) {
            /* calculate final COM */
            res_end = i;
            res_nat = res_end - res_start;
            svmul(1.0/mtot, com, com);

            /* check if COM is outside box */
            copy_rvec(com,new_com);
            switch (unitcell_enum) {
            case euRect:
                put_atoms_in_box(box,1,&new_com);
                break;
            case euTric:
                put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
                break;
            case euCompact:
                put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
                break;
            }
            rvec_sub(new_com,com,shift);
            if (norm2(shift)) {
                if (debug)
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resind+1,
                             res_start+1, res_end+1, PR_VEC(shift));
                for(j=res_start; j<res_end; j++)
                    rvec_inc(x[j],shift);
            }
            clear_rvec(com);
            mtot = 0;

            /* remember start of new residue */
            res_start = i;
        }
        if (i < natoms) {
            /* calc COM */
            m = atom[i].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[i][d];
            mtot += m;

            presnr = atom[i].resind;
        }
    }
}

static void center_x(int ecenter,rvec x[],matrix box,int n,int nc,atom_id ci[]){
    int i,m,ai;
    rvec cmin,cmax,box_center,dx;

    if (nc > 0) {
        copy_rvec(x[ci[0]],cmin);
        copy_rvec(x[ci[0]],cmax);
        for(i=0; i<nc; i++) {
            ai=ci[i];
            for(m=0; m<DIM; m++) {
                if (x[ai][m] < cmin[m])
                    cmin[m]=x[ai][m];
                else if (x[ai][m] > cmax[m])
                    cmax[m]=x[ai][m];
            }
        }
        calc_box_center(ecenter,box,box_center);
        for(m=0; m<DIM; m++)
            dx[m] = box_center[m]-(cmin[m]+cmax[m])*0.5;

        for(i=0; i<n; i++)
            rvec_inc(x[i],dx);
    }
}

static void acc_x(int ncent, int *cindex, rvec *x_in, rvec *x_out){
	for(int ii=0,i=cindex[ii];ii<ncent;ii++,i=cindex[ii]){
		x_out[i][XX]+=x_in[i][XX];
		x_out[i][YY]+=x_in[i][YY];
		x_out[i][ZZ]+=x_in[i][ZZ];
	}
}
static void avg_x(int ncent, int *cindex,int n, rvec *x, rvec  add){
	for(int ii=0,i=cindex[ii];ii<ncent;ii++, i=cindex[ii]){
		x[i][XX]/=(real) n;
		x[i][YY]/=(real) n;
		x[i][ZZ]/=(real) n;
		x[i][XX]+=add[XX];
		x[i][YY]+=add[YY];
		x[i][ZZ]+=add[ZZ];
	}
}
static int calc_dipole(rvec com,rvec dip, int ncent, int * cindex, real *q, rvec *x){
	com[XX]=0.0;com[YY]=0.0;com[ZZ]=0.0;
	dip[XX]=0.0;dip[YY]=0.0;dip[ZZ]=0.0;
	for(int ii=0, i=cindex[ii];ii<ncent;ii++, i=cindex[ii]){
		com[XX]+=x[i][XX];
		com[YY]+=x[i][YY];
		com[ZZ]+=x[i][ZZ];
		dip[XX]+=q[i]*x[i][XX];
		dip[YY]+=q[i]*x[i][YY];
		dip[ZZ]+=q[i]*x[i][ZZ];
	}
	com[XX]/=(real) ncent;
	com[YY]/=(real) ncent;
	com[ZZ]/=(real) ncent;
	return 0;
}

void gmx_meancharges(int argc,char *argv[])
{
  const char *desc[] = {
    "this is a small test program meant to serve as a template ",
    "when writing your own analysis tools. The advantage of ",
    "using gromacs for this is that you have access to all ",
    "information in the topology, and your program will be ",
    "able to handle all types of coordinates and trajectory ",
    "files supported by gromacs. Go ahead and try it! ",
    "This test version just writes the coordinates of an ",
    "arbitrary atom to standard out for each frame. You can ",
    "select which atom you want to examine with the -n argument."
  };
  FILE *fp=NULL;
  FILE *fpxplor=NULL;
  t_topology top;
  int        ePBC;
  char       title[STRLEN], *grpnm;
  t_trxframe fr;
  rvec       *xtop;
  matrix     box;
  int        status,status_out=0;
  int        flags = TRX_READ_X;
  int        flagsb;
  int ncent;
  output_env_t oenv;
  int nx,ny,nz;
  bool bIndex, bCenter=TRUE;
  atom_id      *cindex;
  bool         bPBCWhole,bPBCcomRes,bPBCcomMol,bPBCcomAtom,bPBC,bNoJump,bCluster,bXplor;
  bool 		bOXT=FALSE,bDipole=FALSE;
  bool 		eFilter=FALSE;

  const char  *warn;
  const char     *distfile;
  t_atoms      *atoms=NULL;
  rvec *x;

  int pbc_enum;
  enum
  {
      epSel,
      epNone,
      epComMol,
      epComRes,
      epComAtom,
      epNojump,
      epCluster,
      epWhole,
      epNR
  };
  const char *pbc_opt[epNR + 1] =
      { NULL, "none", "mol", "res", "atom", "nojump", "cluster", "whole",
          NULL };

  int unitcell_enum;
  const char *unitcell_opt[euNR+1] =
      { NULL, "rect", "tric", "compact", NULL };

  enum
  { ecSel, ecTric, ecRect, ecZero, ecNR};
  const char *center_opt[ecNR+1] =
      { NULL, "tric", "rect", "zero", NULL };
  int ecenter;

  t_pargs
		  pa[] ={
				  {"-grid", FALSE, etRVEC, {&ngrid}, "Grid size"},
                  { "-pbc", FALSE, etENUM,
                      { pbc_opt },
                      "PBC treatment (see help text for full description)" },
                  { "-ur", FALSE, etENUM,
                      { unitcell_opt }, "Unit-cell representation" },
                  { "-dipole", FALSE, etBOOL,
                      { &bDipole }, "The center molecule is a dipole" },
                  { "-xplor", FALSE, etBOOL,
                      {&bXplor}, "Write a potential map in Xplor format" },
                  { "-boxcenter", FALSE, etENUM,
                      { center_opt }, "Center for -pbc and -center" },
  };

  t_filenm fnm[] = {
    { efTPS,  NULL,  NULL, ffREAD },   /* this is for the topology */
    { efTRX, "-f", NULL, ffREAD },      /* and this for the trajectory */
    { efXVG, "-dist", "distr",    ffOPTWR  },
    { efTRX, "-oxt","coord.xtc", ffOPTWR },
    { efNDX, NULL,  NULL,         ffOPTRD },
  };
#define NFILE asize(fnm)


  CopyRight(stderr,argv[0]);


  /* This is the routine responsible for adding default options,
   * calling the X/motif interface, etc. */
  parse_common_args(&argc,argv,PCA_CAN_TIME | PCA_CAN_VIEW,
		    NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL, &oenv);

  /* We don't need any topology information to write the coordinates,
   * but to show how it works we start by writing the name and
   * charge of the selected atom. It returns a boolean telling us
   * whether the topology was found and could be read
   */


    pbc_enum   = nenum(pbc_opt);
    bPBCWhole  = pbc_enum==epWhole;
    bPBCcomRes = pbc_enum==epComRes;
    bPBCcomMol = pbc_enum==epComMol;
    bPBCcomAtom= pbc_enum==epComAtom;
    bNoJump    = pbc_enum==epNojump;
    bCluster   = pbc_enum==epCluster;
    bPBC       = pbc_enum!=epNone;
    unitcell_enum = nenum(unitcell_opt);
    ecenter    = nenum(center_opt) - ecTric;

    /*
    int o;
    for(o=0;o<5;o++) printf("%s ",center_opt[o]);
    printf(" %d \n", nenum2(center_opt));
    printf(" ecenters %d %d \n",ecenter,ecTric);
    return;
     */
    distfile = opt2fn_null("-dist",NFILE,fnm);
    bOXT = opt2bSet("-oxt",NFILE,fnm);

    if(distfile) fp = xvgropen(distfile,"","","",oenv);

	if (bOXT) {
		flagsb = flags | TRX_READ_X;
		status_out = open_trx(opt2fn("-oxt",NFILE,fnm),"w");
	}


    nx=(int) *ngrid;
    ny=(int) *(ngrid+1);
    nz=(int) *(ngrid+2);

    Rho_sol = mk_rgrid(nx, ny, nz);
    Rho_slt = mk_rgrid(nx, ny, nz);
    Rho_tot = mk_rgrid(nx, ny, nz);
	Phi_sol = mk_rgrid(nx, ny, nz);
	Phi_slt = mk_rgrid(nx, ny, nz);
	Phi_tot = mk_rgrid(nx, ny, nz);

	int i;

	real *ra=&Rho_sol[0][0][0];
	real *rb=&Rho_slt[0][0][0];
	real *rc=&Rho_tot[0][0][0];

	for (i = 0; i < nx*ny*nz; i++) {
		ra[i] = 0.0;
		rb[i] = 0.0;
		rc[i] = 0.0;
	}

	bIndex = ftp2bSet(efNDX, NFILE, fnm);
//	if (bIndex)
//		printf(" Success !!\n");

	read_tps_conf(ftp2fn(efTPS, NFILE, fnm), title, &top, &ePBC, &xtop, NULL,
			box, TRUE);
	sfree(xtop);
	atoms = &top.atoms;

	if(bOXT){
		int natom=atoms->nr;
		x=calloc(natom,sizeof(rvec));

	}
	if (bCenter) {
		if (bIndex) {
			printf("Select group for centering\n");
			rd_index(ftp2fn_null(efNDX, NFILE, fnm), 1, &ncent, &cindex, &grpnm);
		} else {
			gmx_fatal(FARGS, "Cannot do centering without Index file. Abort.");
		}
	}

//	printf(" %d\n",top.atoms.nr);
	Charges=malloc(sizeof(real)*top.atoms.nr);

	for(i=0;i<top.atoms.nr;i++) {
		Charges[i]=top.atoms.atom[i].q;
	}
  /* The first time we read data is a little special */
	read_first_frame(oenv, &status,ftp2fn(efTRX,NFILE,fnm),&fr,flags);
	hist_t rsol, rslt, psol, pslt,ptot;

	if(setup_rdf_grid(&rsol)) return;
	if(setup_rdf_grid(&rslt)) return;
	if(setup_rdf_grid(&psol)) return;
	if(setup_rdf_grid(&pslt)) return;
	if(setup_rdf_grid(&ptot)) return;

	/* This is the main loop over frames */
	int n=0;
	real Volume=0.0;
	matrix oc, *co;
	rvec com, dip, dip_n;

	do {
		/* coordinates are available in the vector fr.x
		 * you can find this and all other structures in
		 * the types directory under the gromacs include dir.
		 * Note how flags determines whether to read x/v/f!
		 */

		if (bCenter) center_x(ecenter,fr.x,fr.box,fr.natoms,ncent,cindex);
		t_pbc pbc;

		if (bPBCcomAtom) {
			switch (unitcell_enum) {
				case euRect:
				put_atoms_in_box(fr.box,fr.natoms,fr.x);
				break;
				case euTric:
				put_atoms_in_triclinic_unitcell(ecenter,fr.box,fr.natoms,fr.x);
				break;
				case euCompact:
				warn = put_atoms_in_compact_unitcell(ePBC,ecenter,fr.box,
						fr.natoms,fr.x);
				break;
			}
		}
		if (bPBCcomRes) {
			put_residue_com_in_box(unitcell_enum,ecenter,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}
		if (bPBCcomMol) {
			put_molecule_com_in_box(unitcell_enum,ecenter,
					&top.mols,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}

		set_pbc(&pbc,ePBC,fr.box);
		co=&fr.box;

		calc_recipbox(*co, oc);

		if(bDipole){
			calc_dipole(com, dip, ncent,cindex, Charges, fr.x);
		}else gmx_calc_com_pbc(&top,fr.x,&pbc,ncent,cindex,com);


		if (bOXT) {
			acc_x(ncent, cindex, fr.x, x);
		}
		if(GetCharges(Rho_slt,Rho_sol,Charges,nx,ny,nz,ncent,cindex, &fr)) gmx_fatal(FARGS,"Cannot average charges. Abort!\n");

		Volume+=det(fr.box);
		n++;

	} while(read_next_frame(oenv,status,&fr));

	for(i=0;i<nx*ny*nz;i++){
		ra[i]=ra[i]/(real) n;
		rb[i]=rb[i]/(real) n;
	}
	if(distfile) {
		do_rdf_grid(Rho_sol,nx,ny,nz,*co,oc,com,&rsol);
		print_rdf_grid(fp,distfile,rsol,oenv);
		fprintf(fp,"&\n");
	}
	if(distfile) {
		do_rdf_grid(Rho_slt,nx,ny,nz,*co,oc,com,&rslt);
		print_rdf_grid(fp,distfile,rslt,oenv);
		fprintf(fp,"&\n");
	}

	eFilter=TRUE;
	macro_filter(nx,ny,nz, fr.box, 0.4);
	Volume=Volume/(real) n;
	get_potential(Rho_sol, oc, nx,ny,nz, Volume, Phi_sol, eFilter);
	if(distfile) {
		if(bDipole) {
			do_rdf_dip_grid(Phi_sol,nx,ny,nz,*co,oc,com,dip,&psol);
		} else do_rdf_grid(Phi_sol,nx,ny,nz,*co,oc,com,&psol);


		print_rdf_grid(fp,distfile,psol,oenv);
		fprintf(fp,"&\n");
	}
	eFilter=FALSE;
	get_potential(Rho_slt, oc, nx,ny,nz, Volume, Phi_slt, eFilter);
	if(distfile) {
		if(bDipole) {
			do_rdf_dip_grid(Phi_slt,nx,ny,nz,*co,oc,com,dip,&pslt);
		} else 	do_rdf_grid(Phi_slt,nx,ny,nz,*co,oc,com,&pslt);
		print_rdf_grid(fp,distfile,pslt,oenv);
		fprintf(fp,"&\n");
	}

	ra=&Phi_sol[0][0][0];
	rb=&Phi_slt[0][0][0];
	rc=&Phi_tot[0][0][0];
	for(int i=0;i<nx*ny*nz;i++){
		rc[i]=ra[i]+rb[i];
	}

	if(distfile) {
		if(bDipole) {
			do_rdf_dip_grid(Phi_tot,nx,ny,nz,*co,oc,com,dip,&ptot);
		} else do_rdf_grid(Phi_tot,nx,ny,nz,*co,oc,com,&ptot);
		print_rdf_grid(fp,distfile,ptot,oenv);
	}


	if(bXplor) {
		fpxplor=fopen("slt_out.xplor","w+");
		get_xplor(Phi_slt,nx,ny,nz,*co,fpxplor);
		ffclose(fpxplor);
	}
	if(bXplor) {
		fpxplor=fopen("sol_out.xplor","w+");
		get_xplor(Phi_sol,nx,ny,nz,*co,fpxplor);
		ffclose(fpxplor);
	}
	if(bXplor) {
		fpxplor=fopen("tot_out.xplor","w+");
		get_xplor(Phi_tot,nx,ny,nz,*co,fpxplor);
		ffclose(fpxplor);
	}


	if (bOXT) {
		rvec add;
		add[XX]=0.0;
		add[YY]=0.0;
		add[ZZ]=0.0;
		avg_x(ncent, cindex, n, x, add);
		write_trx(status_out,ncent,cindex, &top.atoms,
              1,1.0,*co,x,NULL, NULL);
      }

	ffclose(fp);

	thanx(stderr);

}

